import cv2
import numpy as np
from multiprocessing import Queue
from typing import Any, Dict, List, Tuple


def showImages(left_queue: Queue, right_queue: Queue):
    """显示两个相机的图像"""
    while True:
        if not left_queue.empty() and not right_queue.empty():
            frame_left = left_queue.get()
            frame_right = right_queue.get()

            combined_frame = np.hstack((frame_left, frame_right))

            cv2.imshow("Camera Images", combined_frame)

            if cv2.waitKey(1) & 0xFF == ord("q"):
                break
        else:
            # 等待新的图像帧
            cv2.waitKey(10)

    cv2.destroyAllWindows()


def showUndistortImages(
    left_map1: np.ndarray,
    left_map2: np.ndarray,
    right_map1: np.ndarray,
    right_map2: np.ndarray,
    left_queue: Queue,
    right_queue: Queue,
) -> None:
    """显示校正后的图像"""
    while True:
        if not left_queue.empty() and not right_queue.empty():
            frame_left = left_queue.get()
            frame_right = right_queue.get()

            # 使用预计算的映射矩阵进行校正
            dst_left = cv2.remap(frame_left, left_map1, left_map2, cv2.INTER_LINEAR)
            dst_right = cv2.remap(frame_right, right_map1, right_map2, cv2.INTER_LINEAR)

            # 在校正后的图像上画水平线
            line_color = (0, 255, 0)  # 绿色
            line_thickness = 1
            for y in range(0, dst_left.shape[0], 50):  # 每50行画一条线
                cv2.line(
                    dst_left, (0, y), (dst_left.shape[1], y), line_color, line_thickness
                )
                cv2.line(
                    dst_right,
                    (0, y),
                    (dst_right.shape[1], y),
                    line_color,
                    line_thickness,
                )

            # 可视化检查
            combined_frame = np.hstack((dst_left, dst_right))

            # 显示校正前后的图像对比
            original_combined = np.hstack((frame_left, frame_right))
            full_combined = np.vstack((original_combined, combined_frame))

            cv2.imshow("Original and Corrected Images", full_combined)

            if cv2.waitKey(1) & 0xFF == ord("q"):
                break
        else:
            # 等待新的图像帧
            cv2.waitKey(10)

    cv2.destroyAllWindows()


def showDisparity(
    left_queue: Queue,
    right_queue: Queue,
    left_map1: np.ndarray,
    left_map2: np.ndarray,
    right_map1: np.ndarray,
    right_map2: np.ndarray,
):
    """计算并显示视差图"""

    def updateParameters(val):
        nonlocal num, blockSize
        num = cv2.getTrackbarPos("num", "Disparity Map")
        blockSize = cv2.getTrackbarPos("blockSize", "Disparity Map")
        if blockSize % 2 == 0:
            blockSize += 1  # blockSize 必须为奇数
        if blockSize < 5:
            blockSize = 5  # blockSize 最小为5

    cv2.namedWindow("Disparity Map")
    num = 3
    blockSize = 5
    cv2.createTrackbar("num", "Disparity Map", num, 20, updateParameters)  # 最大为20
    cv2.createTrackbar(
        "blockSize", "Disparity Map", blockSize, 50, updateParameters
    )  # 最大为50

    # 显示一个空的初始图像，确保窗口和滑动条已经完全初始化
    dummy_img = np.zeros((480, 640), dtype=np.uint8)
    cv2.imshow("Disparity Map", dummy_img)
    cv2.waitKey(100)  # 等待100ms，确保滑动条初始化完成

    updateParameters(0)  # 初始化参数

    while True:
        if not left_queue.empty() and not right_queue.empty():
            frame_left = left_queue.get()
            frame_right = right_queue.get()

            rectified_left = cv2.remap(
                frame_left, left_map1, left_map2, cv2.INTER_LINEAR
            )
            rectified_right = cv2.remap(
                frame_right, right_map1, right_map2, cv2.INTER_LINEAR
            )

            gray_left = cv2.cvtColor(rectified_left, cv2.COLOR_BGR2GRAY)
            gray_right = cv2.cvtColor(rectified_right, cv2.COLOR_BGR2GRAY)

            num_disp = 16 * num  # 更新视差范围

            stereo = cv2.StereoSGBM_create(
                minDisparity=0,  # 最小视差，通常为0
                numDisparities=num_disp,  # 视差范围，必须是16的倍数
                blockSize=blockSize,  # 匹配块的大小，通常为5到21之间的奇数
                P1=8
                * 3
                * blockSize**2,  # 控制视差平滑度的第一个参数，值越大，视差图越平滑
                P2=32
                * 3
                * blockSize**2,  # 控制视差平滑度的第二个参数，值越大，视差图越平滑
                disp12MaxDiff=1,  # 左右视差检查的最大容许差异，通常设置为1
                preFilterCap=31,  # 预滤波器的最大值，一般设置为31或63
                uniquenessRatio=10,  # 视差唯一性比例，越大则唯一性约束越强，一般设置为5到15
                speckleWindowSize=100,  # 视差斑点检查的窗口大小，通常设置为50到200
                speckleRange=32,  # 视差斑点检查的最大差异，一般设置为1到32
                mode=cv2.STEREO_SGBM_MODE_SGBM,  # 视差计算模式，常用模式为STEREO_SGBM_MODE_SGBM或STEREO_SGBM_MODE_SGBM_3WAY
            )

            disparity = stereo.compute(gray_left, gray_right).astype(np.float32) / 16.0

            # 生成彩色视差图
            disparity_color = cv2.applyColorMap(
                cv2.normalize(disparity, None, 0, 255, cv2.NORM_MINMAX).astype(
                    np.uint8
                ),
                cv2.COLORMAP_JET,
            )

            cv2.imshow("Disparity Map", disparity_color)

            if cv2.waitKey(1) & 0xFF == ord("q"):
                break
        else:
            # 等待新的图像帧
            cv2.waitKey(10)

    cv2.destroyAllWindows()
